#include "mnist.h"
#include <conio.h> // 为了使用 _getch()

int main() {
    float *train_images;
    unsigned char *train_labels;
    int num_train_items;
    
    float *test_images;
    unsigned char *test_labels;
    int num_test_items;
    
    printf("Loading training data...\n");
    read_mnist_data("train-images.idx3-ubyte", 
                   "train-labels.idx1-ubyte",
                   &train_images, &train_labels, &num_train_items);
    printf("Successfully loaded %d training images\n", num_train_items);
                   
    printf("Loading test data...\n");
    read_mnist_data("t10k-images.idx3-ubyte", 
                   "t10k-labels.idx1-ubyte",
                   &test_images, &test_labels, &num_test_items);
    printf("Successfully loaded %d test images\n", num_test_items);
    
    printf("Creating neural network...\n");
    Network *net = create_network();
    printf("Starting training...\n");
    train_network(net, train_images, train_labels, num_train_items);
    
    printf("Testing network...\n");
    int correct = 0;
    for(int i = 0; i < num_test_items; i++) {
        int prediction = predict(net, &test_images[i * IMAGE_SIZE]);
        if(prediction == test_labels[i]) {
            correct++;
        }
        if((i + 1) % 1000 == 0) {
            printf("Tested %d images...\n", i + 1);
        }
    }
    
    printf("Testing completed!\n");
    printf("Accuracy: %.2f%%\n", 100.0f * correct / num_test_items);
    
    // 添加交互式测试循环
    char continue_test = 'y';
    while(continue_test == 'y' || continue_test == 'Y') {
        printf("\nEnter test image number (0-%d): ", num_test_items - 1);
        int test_index;
        scanf("%d", &test_index);
        
        if(test_index >= 0 && test_index < num_test_items) {
            float *test_image = &test_images[test_index * IMAGE_SIZE];
            print_image(test_image);
            print_probabilities(net, test_image);
            printf("Actual digit: %d\n", test_labels[test_index]);
            
            printf("\nPress any key to continue...");
            _getch(); // 等待用户按键
            
            printf("\nTest another image? (y/n): ");
            scanf(" %c", &continue_test);
        } else {
            printf("Invalid image number\n");
        }
    }
    
    // 释放内存
    free_network(net);
    free(train_images);
    free(train_labels);
    free(test_images);
    free(test_labels);
    
    printf("\nAll resources have been freed successfully!\n");
    printf("Press any key to exit...");
    _getch();
    
    return 0;
}
